package speech.tag;

import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.Reader;
import java.util.List;

import be.ac.ulg.montefiore.run.jahmm.Hmm;
import be.ac.ulg.montefiore.run.jahmm.ObservationInteger;
import be.ac.ulg.montefiore.run.jahmm.OpdfIntegerFactory;
import be.ac.ulg.montefiore.run.jahmm.io.FileFormatException;
import be.ac.ulg.montefiore.run.jahmm.io.ObservationIntegerReader;
import be.ac.ulg.montefiore.run.jahmm.io.ObservationSequencesReader;
import be.ac.ulg.montefiore.run.jahmm.learn.BaumWelchLearner;
import be.ac.ulg.montefiore.run.jahmm.toolbox.KullbackLeiblerDistanceCalculator;
import speech.hmm.LeftRightHmm;

public class TagWrapper {
	private String tag;									// tag
	private Hmm<ObservationInteger> learntHmm = null;	// trained HMM
	private int noStates;								// number of states of the HMM
	private int dictionarySize;							// dictionary size, used for the initial opdf
	private int hmmDelta;								// HMM delta
	private String trainingFile;						// training file path
	
	public TagWrapper(String tag, int noStates, int dictionarySize, int hmmDelta) {
		this.tag = tag;
		this.noStates = noStates;
		this.dictionarySize = dictionarySize;
		this.hmmDelta = hmmDelta;
	}
	
	public TagWrapper(String tag, int noStates, int dictionarySize, String trainingFile, int hmmDelta) {
		this(tag, noStates, dictionarySize, hmmDelta);
		this.trainingFile = trainingFile;
	}
	
	/**
	 * Read sequences in the training file. 
	 * 
	 * @return
	 * @throws IOException
	 * @throws FileFormatException
	 */
	private List<List<ObservationInteger>> readTrainingSet(String filepath) throws IOException, FileFormatException {
		Reader reader = new FileReader(filepath);
		List<List<ObservationInteger>> sequences = ObservationSequencesReader.readSequences(new ObservationIntegerReader(), reader);
		reader.close();
		
		return sequences;
	}
	
	/**
	 * Train HMM on the training file specified by filepath.
	 * 
	 * @param filepath
	 */
	public void trainModel(String filepath) {
		List<List<ObservationInteger>> sequences = null;
		
		try {
			sequences = readTrainingSet(filepath);
		} catch (Exception e) {
			e.printStackTrace();
		}
		
		KullbackLeiblerDistanceCalculator klc = new KullbackLeiblerDistanceCalculator();
		Hmm<ObservationInteger> cloneHmm;
		
		try {
			LeftRightHmm<ObservationInteger> hmm = new LeftRightHmm<ObservationInteger>(noStates,
					new OpdfIntegerFactory(dictionarySize), hmmDelta);
			cloneHmm = hmm.clone();
			
			BaumWelchLearner bwl = new BaumWelchLearner();
			learntHmm = bwl.learn(cloneHmm, sequences);
			
			System.out.println("Distance after learning: " + klc.distance(learntHmm, cloneHmm));
		} catch (CloneNotSupportedException e) {
			e.printStackTrace();
		}
	}
	
	public void trainModel() {
		trainModel(trainingFile);
	}
	
	/**
	 * Calculate the probability of a sequence on the HMM.
	 * 
	 * @param seq
	 * @return
	 */
	public double getProbability(List<ObservationInteger> seq) {
		if (learntHmm == null) return -1;
		
		return learntHmm.probability(seq);
	}
	
	public String getTag() {
		return tag;
	}

	public String toString() {
		if (learntHmm != null) return learntHmm.toString();
		
		return null;
	}
}
